# -*- coding: utf-8 -*-

# backend/utils/websocket_metrics.py
from typing import Dict, Any, Optional
from datetime import datetime, timedelta
from collections import defaultdict, deque
import asyncio


class WebSocketMetrics:
    """WebSocket指标收集器"""
    
    def __init__(self):
        self.connection_count = 0
        self.total_connections = 0
        self.message_count = 0
        self.error_count = 0
        
        # 按类型统计消息
        self.message_types: Dict[str, int] = defaultdict(int)
        
        # 按用户统计
        self.user_connections: Dict[str, int] = defaultdict(int)
        self.user_messages: Dict[str, int] = defaultdict(int)
        
        # 性能指标
        self.response_times: deque = deque(maxlen=1000)
        self.message_sizes: deque = deque(maxlen=1000)
        
        # 时间序列数据
        self.hourly_stats: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))
        
        # 启动清理任务
        asyncio.create_task(self._cleanup_task())
    
    def record_connection(self, user_id: str):
        """记录新连接"""
        self.connection_count += 1
        self.total_connections += 1
        self.user_connections[user_id] += 1
    
    def record_disconnection(self, user_id: str):
        """记录断开连接"""
        self.connection_count = max(0, self.connection_count - 1)
        self.user_connections[user_id] = max(0, self.user_connections[user_id] - 1)
    
    def record_message(self, message_type: str, user_id: str, message_size: int):
        """记录消息"""
        self.message_count += 1
        self.message_types[message_type] += 1
        self.user_messages[user_id] += 1
        self.message_sizes.append(message_size)
        
        # 记录小时统计
        hour_key = datetime.utcnow().strftime("%Y-%m-%d-%H")
        self.hourly_stats[hour_key]["messages"] += 1
        self.hourly_stats[hour_key][f"type_{message_type}"] += 1
    
    def record_error(self, error_type: str):
        """记录错误"""
        self.error_count += 1
        
        hour_key = datetime.utcnow().strftime("%Y-%m-%d-%H")
        self.hourly_stats[hour_key]["errors"] += 1
        self.hourly_stats[hour_key][f"error_{error_type}"] += 1
    
    def record_response_time(self, response_time: float):
        """记录响应时间"""
        self.response_times.append(response_time)
    
    def get_metrics(self) -> Dict[str, Any]:
        """获取指标数据"""
        
        # 计算平均响应时间
        avg_response_time = 0
        if self.response_times:
            avg_response_time = sum(self.response_times) / len(self.response_times)
        
        # 计算平均消息大小
        avg_message_size = 0
        if self.message_sizes:
            avg_message_size = sum(self.message_sizes) / len(self.message_sizes)
        
        return {
            "connections": {
                "current": self.connection_count,
                "total": self.total_connections,
                "by_user": dict(self.user_connections)
            },
            "messages": {
                "total": self.message_count,
                "by_type": dict(self.message_types),
                "by_user": dict(self.user_messages),
                "avg_size": avg_message_size
            },
            "performance": {
                "avg_response_time": avg_response_time,
                "error_count": self.error_count,
                "error_rate": self.error_count / max(1, self.message_count) * 100
            },
            "hourly_stats": dict(self.hourly_stats)
        }
    
    def get_connection_stats(self) -> Dict[str, Any]:
        """获取连接统计"""
        return {
            "active_connections": self.connection_count,
            "total_connections": self.total_connections,
            "unique_users": len([u for u, c in self.user_connections.items() if c > 0]),
            "avg_connections_per_user": self.connection_count / max(1, len(self.user_connections))
        }
    
    async def _cleanup_task(self):
        """清理过期数据"""
        while True:
            try:
                await asyncio.sleep(3600)  # 每小时清理一次
                
                # 清理24小时前的数据
                cutoff_time = datetime.utcnow() - timedelta(hours=24)
                cutoff_key = cutoff_time.strftime("%Y-%m-%d-%H")
                
                # 移除过期的小时统计
                expired_keys = [k for k in self.hourly_stats.keys() if k < cutoff_key]
                for key in expired_keys:
                    del self.hourly_stats[key]
                
            except Exception as e:
                print(f"清理任务错误: {str(e)}")


# 全局指标实例
websocket_metrics = WebSocketMetrics()
